import { NextRequest } from "next/server";
import { ChatOllama } from "@langchain/ollama";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from '@langchain/core/output_parsers'; //流式
import { RunnableLambda, RunnableSequence, RunnableWithMessageHistory, RunnablePassthrough } from '@langchain/core/runnables';
import { JsonOutputParser } from "@langchain/core/output_parsers";
import { InMemoryChatMessageHistory, BaseListChatMessageHistory } from '@langchain/core/chat_history';
import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
import { OllamaEmbeddings } from "@langchain/ollama";
import { MessagesPlaceholder } from '@langchain/core/prompts';
import { Document } from "@langchain/core/documents";
export const runtime = "nodejs";

// 初始化向量存储
let vectorStore: PGVectorStore | null = null;

async function initVectorStore() {
  if (!vectorStore) {
    vectorStore = await PGVectorStore.initialize(
      new OllamaEmbeddings({
        baseUrl: "http://192.168.31.159:11434",
        model: "quentinz/bge-large-zh-v1.5:latest",
      }),
      {
        tableName: "embedding_zqw_echart_1024",
        columns: {
          idColumnName: "embedding_id",
          vectorColumnName: "embedding",
          contentColumnName: "text",
          metadataColumnName: "metadata",
        },
        postgresConnectionOptions: {
          host: "118.178.106.254",
          port: 15432,
          user: "postgres",
          password: "6a29c7071ae1f1b9",
          database: "postgres",
        },
      }
    );
  }
  return vectorStore;
}
const chatHistoryCache = new Map<string, InMemoryChatMessageHistory>();
// 创建 Ollama 实例
const ollama = new ChatOllama({
  baseUrl: 'http://192.168.31.159:11434',
  model: 'qwen:32b', // 使用您想要的模型
  dangerouslyAllowBrowser: true,
});


// 图表配置的JSON Schema
const chartConfigSchema = {
  type: "object",
  properties: {
    title: {
      type: "object",
      properties: {
        text: { type: "string" },
        left: { type: "string" },
      },
    },
    tooltip: {
      type: "object",
      properties: {
        trigger: { type: "string" },
      },
    },
    legend: {
      type: "object",
      properties: {
        data: { type: "array", items: { type: "string" } },
        top: { type: "string" },
      },
    },
    xAxis: {
      type: "object",
      properties: {
        type: { type: "string" },
        data: { type: "array", items: { type: "string" } },
      },
    },
    yAxis: {
      type: "object",
      properties: {
        type: { type: "string" },
      },
    },
    series: {
      type: "array",
      items: {
        type: "object",
        properties: {
          name: { type: "string" },
          type: { type: "string" },
          data: { type: "array" },
        },
      },
    },
  },
};

const TEMPLATE = `你是一个专业的数据分析师，擅长从表格数据中提取信息并生成ECharts图表配置。

你的任务：
1. 分析用户提供的表格数据
2. 理解数据的含义和结构
3. 生成适合的ECharts图表配置（JSON格式）

图表类型选择规则：
- 如果数据是分类对比（如不同类别的数值对比），使用柱状图（bar）
- 如果数据是趋势变化（如时间序列），使用折线图（line）
- 如果数据是占比关系，使用饼图（pie）
- 如果数据是相关性分析，使用散点图（scatter）

输出要求：
- 必须输出有效的JSON格式
- JSON必须符合ECharts配置格式
- 包含title、tooltip、legend、xAxis、yAxis、series等必要字段
- 图表标题要清晰描述数据内容
- 数据要准确对应表格内容

输出格式示例：
{{
  "title": {{
    "text": "图表标题",
    "left": "center"
  }},
  "tooltip": {{
    "trigger": "axis"
  }},
  "legend": {{
    "data": ["系列1", "系列2"],
    "top": "10%"
  }},
  "xAxis": {{
    "type": "category",
    "data": ["类别1", "类别2", "类别3"]
  }},
  "yAxis": {{
    "type": "value"
  }},
  "series": [
    {{
      "name": "系列1",
      "type": "bar",
      "data": [10, 20, 30]
    }}
  ]
}}`

export async function POST(req: NextRequest) {
  try {
    const body = await req.json();
    const { question, sessionId } = body;
    if (!question) return Response.json({ ok: false, message: "缺少查询问题" }, { status: 400 });
    console.log("🔍 开始查询:", question);

    // 初始化向量存储
    const vectorStore = await initVectorStore();

    // 检索相关文档
    const retriever = vectorStore.asRetriever(5); // 检索前5个最相关的文档
    const chatPrompt = ChatPromptTemplate.fromMessages([
      ['system', TEMPLATE],
      new MessagesPlaceholder('history'), //hisotory占位符
      ['human', `表格数据：\n{context}\n\n用户问题：{question}`]
    ]);
    // 定义问题改写的chain
    const rephraseChainPrompt = ChatPromptTemplate.fromMessages([
      [
        'system',
        '给定以下对话和一个后续问题，请将后续问题重述为一个独立的问题。请注意，重述的问题应该包含足够的信息，使得没有看过对话历史的人也能理解。',
      ],
      new MessagesPlaceholder('history'),
      ['human', '将以下问题重述为一个独立的问题：\n{question}'],
    ]);
    // 改写提问的chain
    const rephraseChain = RunnableSequence.from([rephraseChainPrompt, ollama, new StringOutputParser()]);
    // 上下文检索Chain
    const contextRetrieveChain = RunnableSequence.from([
      (input: { question: string }) => input.standalone_question, //提取用户问题
      retriever, //检索相关原文
      (docs: Document[]) => {
        console.log(
          '🔍 检索到的文档块：',
          docs.map((doc) => doc.pageContent)
        );
        return docs.map((doc) => doc.pageContent).join('\n');
      },
    ]);
    const ragChain = RunnableSequence.from([
      RunnablePassthrough.assign({
        standalone_question: rephraseChain,
      }),
      RunnablePassthrough.assign({
        context: contextRetrieveChain,
      }),
      chatPrompt, // 注入context和question到Prompt
      ollama, // 调用LLM生成回答
      new JsonOutputParser(), // 把LLM输出转成纯文本
    ]);
    const chainWithHistory = new RunnableWithMessageHistory({
      runnable: ragChain,
      // getMessageHistory接收一个函数 需要传入_sessionId 去获取ChatMessageHistory对象
      // getMessageHistory: (_sessionId) => new LocalStorageChatHistory({ sessionId: _sessionId }),
      getMessageHistory: (sessionId: string) => {
        if (!chatHistoryCache.has(sessionId)) {
          // 不存在则新建，并存入缓存
          chatHistoryCache.set(sessionId, new InMemoryChatMessageHistory());
        }
        // 存在则直接返回缓存的实例（复用同一个）
        return chatHistoryCache.get(sessionId)!;
      },
      //以什么key传递消息
      inputMessagesKey: 'question',
      // historyMessagesKey，聊天记录在 prompt 中的 key，因为要自动的把聊天记录注入到 prompt 中。
      historyMessagesKey: 'history',
    });

    const result = await chainWithHistory.invoke({ question }, { configurable: { sessionId } });

    return Response.json({
      ok: true,
      data: {
        chartConfig: result,
        note: "已使用历史对话上下文生成图表",
      },
    });
  } catch (error: unknown) {
    const errorMessage = error instanceof Error ? error.message : String(error);
    console.error("❌ 查询失败:", error);
    return Response.json(
      {
        ok: false,
        message: `查询失败: ${errorMessage}`,
      },
      { status: 500 }
    );
  }
}

